"""
Action Classification Module

Automatically identifies action types by comparing segments
against template library using feature similarity.
"""

import numpy as np
from typing import Dict, List, Tuple, Optional
import logging

from .matcher import SoftDTWMatcher

logger = logging.getLogger(__name__)


class ActionClassifier:
    """
    Classifies action segments by matching against template library.

    Uses DINO feature similarity to identify which action type
    a segment most likely represents.
    """

    def __init__(self,
                 confidence_threshold: float = 0.3,
                 gamma: float = 0.1):
        """
        Initialize action classifier.

        Args:
            confidence_threshold: Minimum similarity score to consider valid match
            gamma: Soft-DTW smoothing parameter
        """
        self.confidence_threshold = confidence_threshold
        self.softdtw_matcher = SoftDTWMatcher(gamma=gamma)

    def classify_segment(self,
                        segment_features: np.ndarray,
                        templates: Dict[str, Dict]) -> Tuple[str, float]:
        """
        Classify a single action segment.

        Args:
            segment_features: DINO features for the segment (T, 768)
            templates: Dictionary of templates {action_type: template_data}

        Returns:
            Tuple of (action_type, confidence_score)
            Returns ('unknown', 0.0) if no confident match found
        """
        if len(templates) == 0:
            logger.warning("No templates available for classification")
            return 'unknown', 0.0

        best_action = None
        best_similarity = -np.inf

        # Compare with all templates
        for action_type, template in templates.items():
            # Compute similarity using Soft-DTW distance
            dtw_distance = self.softdtw_matcher.align_sequences(
                segment_features,
                template['features']
            )

            # Convert distance to similarity score (0-1 range)
            # Lower distance = higher similarity
            similarity = np.exp(-dtw_distance / 5.0)

            logger.debug(f"Segment vs {action_type}: similarity={similarity:.3f}, dtw_distance={dtw_distance:.3f}")

            if similarity > best_similarity:
                best_similarity = similarity
                best_action = action_type

        # Check confidence threshold
        if best_similarity < self.confidence_threshold:
            logger.warning(f"Low confidence match: {best_action} ({best_similarity:.3f} < {self.confidence_threshold})")
            return 'unknown', best_similarity

        logger.info(f"Classified as '{best_action}' with confidence {best_similarity:.3f}")

        return best_action, best_similarity

    def classify_segments(self,
                         segments_features: List[np.ndarray],
                         templates: Dict[str, Dict]) -> List[Tuple[str, float]]:
        """
        Classify multiple action segments.

        Args:
            segments_features: List of feature arrays for each segment
            templates: Dictionary of templates

        Returns:
            List of (action_type, confidence) tuples
        """
        results = []

        for i, seg_features in enumerate(segments_features):
            action_type, confidence = self.classify_segment(seg_features, templates)
            results.append((action_type, confidence))
            logger.info(f"Segment {i+1}: {action_type} (confidence={confidence:.3f})")

        return results

    def classify_with_landmarks(self,
                               segment_features: np.ndarray,
                               segment_landmarks: List[np.ndarray],
                               templates: Dict[str, Dict]) -> Tuple[str, float, Dict]:
        """
        Classify segment and return both action type and geometric similarity.

        Args:
            segment_features: DINO features for the segment
            segment_landmarks: Landmarks for the segment
            templates: Dictionary of templates

        Returns:
            Tuple of (action_type, feature_confidence, geometric_scores)
        """
        # Feature-based classification
        action_type, confidence = self.classify_segment(segment_features, templates)

        # Additional geometric validation (optional)
        geometric_scores = {}

        if action_type != 'unknown' and action_type in templates:
            template = templates[action_type]

            # Simple geometric check: compare ROM patterns
            if 'landmarks' in template:
                from .skeleton_normalizer import calculate_joint_angles

                # Calculate ROM for segment
                seg_angles = [calculate_joint_angles(lm) for lm in segment_landmarks]
                template_angles = [calculate_joint_angles(lm) for lm in template['landmarks']]

                # Compare ROM ranges
                for joint in ['left_knee', 'right_knee', 'left_hip', 'right_hip']:
                    seg_rom = max([a[joint] for a in seg_angles]) - min([a[joint] for a in seg_angles])
                    tmp_rom = max([a[joint] for a in template_angles]) - min([a[joint] for a in template_angles])

                    if tmp_rom > 0:
                        rom_similarity = min(seg_rom / tmp_rom, tmp_rom / seg_rom)
                        geometric_scores[joint] = rom_similarity

        return action_type, confidence, geometric_scores
